import os.path
import random

import tqdm
from gxl_ai_utils.utils import utils_file


def handle_data():
    old_list = utils_file.load_list_file_clean('./gxl_all_2.list')
    added_list = utils_file.load_list_file_clean(
        "/home/work_nfs14/xlgeng/asr_data_shard/wenetspeech_new_all/shards_list.txt")
    old_list_filter = []
    num4del = 0
    for line in old_list:
        if line.startswith("/home/local_data/data4w/shard_1/train_l/") or line.startswith(
                "/home/local_data/hwang/huawei_cn_en/cn2/wenetspeech/"):
            num4del += 1
            if num4del % 1000 == 0:
                utils_file.logging_print("del:{}".format(num4del))
        else:
            old_list_filter.append(line)
    utils_file.logging_print("num4del: " + str(num4del))
    new_list = added_list + old_list_filter
    random.shuffle(new_list)
    utils_file.write_list_to_file(new_list, "./gxl_all_new_wenetspeech.list")


def handle_kespeech(partion='train'):
    speech_dir = f"/home/work_nfs7/hfxue/data/KeSpeech/{partion}"
    text_path = os.path.join(speech_dir, "text")
    wav_path = os.path.join(speech_dir, "wav.scp")
    language_path = os.path.join(speech_dir, "language")
    text_dict = utils_file.load_dict_from_scp(text_path)
    wav_dict = utils_file.load_dict_from_scp(wav_path)
    laguage_dict = utils_file.load_dict_from_scp(language_path)
    language_kinds = list(set(laguage_dict.values()))
    print(language_kinds)
    # >>> ['Northeastern', 'Jiao-Liao', 'Mandarin', 'Ji-Lu', 'Southwestern', 'Zhongyuan', 'Beijing', 'Lan-Yin', 'Jiang-Huai']
    utils_file.logging_print('首先得到中原的音频')

    for area in language_kinds:
        wav_key_list = [k for k, v in laguage_dict.items() if v == area]
        utils_file.logging_print('{}的音频有{}个'.format(area, len(wav_key_list)))
        output_data_dir = f"./kespeech/{area}/{partion}/"
        utils_file.makedir(output_data_dir)
        text_dict_zy = {}
        wav_path_zy = {}
        for key in wav_key_list:
            """"""
            text_dict_zy[key] = text_dict[key]
            wav_path_zy[key] = wav_dict[key]
        utils_file.write_dict_to_scp(text_dict_zy, os.path.join(output_data_dir, "text"))
        utils_file.write_dict_to_scp(wav_path_zy, os.path.join(output_data_dir, "wav.scp"))


def do_handle_fix_list():
    input_path = './gxl_all_new_wenetspeech.list'
    output_path = './gxl_all_new_wenetspeech_fix.list'
    lines = utils_file.load_list_file_clean(input_path)
    lines_res = []
    for line in tqdm.tqdm(lines, desc="handle data", total=len(lines)):
        line = line.replace('local_data', '41_data')
        lines_res.append(line)
    utils_file.write_list_to_file(lines_res, output_path)


def kespeech():
    """"""
    data_dir = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn/examples/aishell/salmonn/data_list/kespeech"
    datanames = os.listdir(data_dir)
    for dataname in datanames:
        data_dir_i = os.path.join(data_dir, dataname, 'test')
        utils_file.do_convert_wav_text_scp_to_jsonl(os.path.join(data_dir_i, 'wav.scp'),
                                                    os.path.join(data_dir_i, 'text'),
                                                    os.path.join(data_dir_i, 'data.list'))


if __name__ == '__main__':
    """"""
    kespeech()
